import numpy as np
import matplotlib.pyplot as plt


def dft(a, sign, opt):
    # opt为光学傅里叶变换，变换后会将0频率放到放到中央，
    n = len(a)
    # 检查数据是否为复数类型，不是则转变为复数类型
    if a.dtype != np.complex:
        a = a.astype(np.complex)

    if opt:
        # 将0频率放到中央的做法是将奇数位变成负数
        a[1:n + 1:2] = -a[1:n + 1:2]
    # 在0到2π产生那个数最为theta（𝜃）
    theta = np.linspace(0, 2 * np.pi, n, endpoint=False)
    # 求对应的sin cos 值放入表中，因为计算sin 和 cos 比较费时间
    # sign为符号位，因为-sin𝜃 = sin-𝜃 所以乘一个符号位
    # 同时因为-cos𝜃 = cos𝜃 就不乘
    sin_tbl = sign * np.sin(theta)
    cos_tbl = np.cos(theta)

    b = np.zeros(n, dtype=np.complex)
    for i in range(n):
        for j in range(n):
            # 因为超过n之后计算的值会重复 k又属于0 -> n-1
            # 所以这样计算出来下标到计算出来的sin_tbl、cos_tbl中直接取值，可以降低时间复杂度
            k = (i * j) % n
            c = np.complex(cos_tbl[k], sin_tbl[k])
            # 由公式得来的
            b[i] += a[j] * c

    if sign > 0:
        # sign 大于0代表逆变换，还需要除n    小于0就是正变换
        b *= 1.0 / n

    if opt:
        a[1:n + 1:2] = -a[1:n + 1:2]

    return b


def test():
    # a = np.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=np.complex)
    # 这里是numpy中的快速傅里叶变换，默认的是不做光学变换
    # b = np.fft.fft(a)

    # 这里是我们自己写的傅里叶变换，在这里我们要与上面对应的话要将true改为false
    # c = dft(a, -1, True)
    # 精度为打印4 不然会精度会很长
    # np.set_printoptions(precision=4, suppress=True)
    # print(b)
    # print(c)


    # 生成信号，并进行傅里叶变换
    x = np.linspace(0, 100, 1001)
    #a = np.sin(x)
    a = np.zeros(len(x),dtype=float)
    for i in range(len(x)):
        if (i//50) %2 ==0:
            a[i] = 1
    #这里传入true相当于做了，光学的傅里叶变换，所以打印的图像就是0频就在中间，
    # 否则0频就应该在开始的位置
    b = dft(a,-1,True)
    # 显示原始信号，傅里叶振幅谱和相位谱
    ax1 = plt.subplot(311)
    ax1.set_title('src signal')
    plt.plot(x,a,'r-')

    ax2 = plt.subplot(312)
    ax2.set_title('amp spectrum')
    # 这里是对复数的实部和虚部进行平方后开根号的到模值在绘制就是频谱图
    plt.plot(x, np.abs(b), 'b-')

    ax3 = plt.subplot(313)
    ax3.set_title('phase spectrum')
    # 这里显示的是复角
    plt.plot(x, np.angle(b), 'g-')

    plt.show()


test()
